import asyncio
import json
from typing import Dict

from starlette.websockets import WebSocket

from logger import logger

active_connections: Dict[str, WebSocket] = {}


class ConnectionManager:
    def __init__(self):
        ...

    async def connect(self, token, websocket: WebSocket):
        print(f"connection token={token}")
        await websocket.accept()
        active_connections[token] = websocket

    def disconnect_token(self, token):
        ws = active_connections[token]
        ws.close()
        del ws

    def disconnect(self, websocket: WebSocket):
        global active_connections
        active_connections = {k: v for k, v in active_connections.items() if v != websocket}

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

    async def broadcast(self, message: str):
        global active_connections
        for k in active_connections:
            print(f"{message}")
            connection = active_connections[k]
            # print(connection)
            await connection.send_text(message)

    def send_with_token(self,  msg):
        global active_connections
        websocket = active_connections.get(msg.token)
        if websocket is None:
            logger.debug(f"websocket为空 token={msg.token}")
        else:
            msg.token=""
            asyncio.run(websocket.send_text(json.dumps(msg.__dict__)))
